package com.gitee.zhuhjay.xunfei.sdk;

import cn.hutool.core.collection.ListUtil;
import cn.hutool.core.text.StrBuilder;
import cn.hutool.core.thread.ExecutorBuilder;
import cn.hutool.core.util.StrUtil;
import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import com.gitee.zhuhjay.xunfei.sdk.consts.ApiVersion;
import com.gitee.zhuhjay.xunfei.sdk.consts.ConfigConst;
import com.gitee.zhuhjay.xunfei.sdk.consts.Role;
import com.gitee.zhuhjay.xunfei.sdk.model.*;
import com.gitee.zhuhjay.xunfei.sdk.util.XfUtil;
import lombok.extern.slf4j.Slf4j;
import okhttp3.*;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

import java.util.List;
import java.util.concurrent.Callable;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Future;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 讯飞AI交互客户端
 *
 * @author ZhuHJay
 * @since 2023/11/2 9:50
 */
@Slf4j
public class AppClient {

    private final ApiAuth apiAuth;
    private final ApiVersion apiVersion;
    /**
     * 默认模型参数
     */
    private static final ModelParams DEFAULT_PARAMS = ModelParams.builder().build();
    /**
     * 自定义模型参数
     */
    private ModelParams customParams;
    /**
     * 注册 function, 只有在v3.0版本才会生效
     */
    private List<Function> functions;
    private final OkHttpClient okHttpClient;
    private final ExecutorService executorService;
    private final AtomicInteger counter = new AtomicInteger(0);

    /**
     * 客户端构造器
     * @param authConfig      api授权信息
     * @param apiVersion      api版本
     * @param okHttpClient    http客户端
     * @param qps             api调用的并发数限制
     */
    AppClient(ApiAuth authConfig, ApiVersion apiVersion, OkHttpClient okHttpClient, int qps) {
        authConfig.check();
        if (okHttpClient == null) {
            throw new IllegalArgumentException("okHttpClient can not be null");
        }
        XfUtil.checkQps(qps);
        this.apiAuth = authConfig;
        this.apiVersion = apiVersion;
        this.okHttpClient = okHttpClient;
        // 线程池核心线程最多开启10个
        this.executorService = ExecutorBuilder.create()
                .setCorePoolSize(Math.min(qps, 10))
                .setMaxPoolSize(qps)
                .setWorkQueue(new LinkedBlockingQueue<>(128))
                .setThreadFactory(r -> new Thread(r,
                        StrUtil.format("{}:{}:{}", ConfigConst.THREAD_FACTORY_NAME, apiAuth.getAppId(), counter.incrementAndGet())))
                .buildFinalizable();
    }

    /**
     * 发送异步请求
     * <p>Tip: 如果使用了 Function Call, 尽量不使用请求历史</p>
     *
     * @param content 提问内容
     * @param history 请求历史(可选)
     * @return ApiResponse
     */
    public Future<ApiResponse> sendAsync(String content, Text... history) {
        if (StrUtil.isBlank(content)) {
            throw new IllegalArgumentException("content can not be blank");
        }
        // 提交任务，等待结果
        return executorService.submit(new WsListener(content, ListUtil.toList(history)));
    }

    /**
     * 自行设置模型参数(将自行检测参数合法性, 如果想使用默认则设置为null)
     *
     * @param params 模型参数
     */
    public void customModelParams(ModelParams params) {
        if (params == null) {
            this.customParams = null;
            return;
        }
        params.check(apiVersion);
        this.customParams = params;
    }

    /**
     * 自行设置function方法(仅有v3.0版本支持且有效)
     *
     * @param functions function 方法
     */
    public void customFunction(List<Function> functions) {
        if (functions == null) {
            this.functions = null;
            return;
        }
        this.functions = functions;
    }

    private class WsListener extends WebSocketListener implements Callable<ApiResponse> {
        private ApiResponse resp;
        private final String question;
        private final List<Text> history;
        private final StrBuilder sb;

        WsListener(String question, List<Text> history) {
            this.sb = StrBuilder.create();
            this.question = question;
            this.history = history;
        }

        @Override
        public void onMessage(@NotNull WebSocket webSocket, @NotNull String res) {
            log.trace("<== on message == \n\t{}", res);
            JSONObject parseObj = JSONUtil.parseObj(res);
            JSONObject header = parseObj.getJSONObject("header");
            // code 不为0，表示出错
            Integer code = header.getInt("code");
            if (code != 0) {
                resp = ApiResponse.builder()
                        .ok(false)
                        .code(code)
                        .message(header.getStr("message"))
                        .question(question)
                        .build();
                // 抛出异常，中断线程
                throw new ApiRequestException();
            }
            Integer status = header.getInt("status");
            JSONObject payload = parseObj.getJSONObject("payload");
            JSONObject choices = payload.getJSONObject("choices");

            // 获取此次接收到的内容
            JSONObject text = choices.getJSONArray("text").toList(JSONObject.class).get(0);
            String content = text.getStr("content");
            sb.append(content);
            if (status == 2) {
                JSONObject counter = payload.getByPath("usage.text", JSONObject.class);
                resp = ApiResponse.builder()
                        .ok(true)
                        .message("ok")
                        .question(question)
                        .answer(sb.toString())
                        .promptTokens(counter.getInt("prompt_tokens"))
                        .completionTokens(counter.getInt("completion_tokens"))
                        .totalTokens(counter.getInt("total_tokens"))
                        .build();
                // 判断是否命中function
                if (text.containsKey("function_call")) {
                    FunctionCall call = FunctionCall.builder()
                            .arguments(text.getByPath("function_call.arguments", String.class))
                            .name(text.getByPath("function_call.name", String.class))
                            .build();
                    resp.setFunctionCall(call);
                }
                synchronized (this) {
                    // 唤醒返回结果的线程(callable)
                    this.notify();
                }
            }
        }

        @Override
        public void onFailure(@NotNull WebSocket webSocket, @NotNull Throwable t, @Nullable Response response) {
            if (!(t instanceof ApiRequestException)) {
                // 连接错误(key、secret错误等，此时没有完成连接，不会触发onMessage)
                resp = ApiResponse.builder()
                        .ok(false)
                        .code(-1)
                        .message(t.getMessage())
                        .question(question)
                        .build();
            }
            synchronized (this) {
                // 此时必须强制唤醒返回结果的线程(callable)
                this.notify();
            }
        }

        @Override
        public ApiResponse call() throws Exception {
            log.trace("==> ready == \n\t{}", question);
            WebSocket ws = okHttpClient.newWebSocket(
                    new Request.Builder()
                            .url(XfUtil.genAuthUrl(apiAuth, apiVersion))
                            .build(), this);
            log.trace("==> connection server == success");
            Text text = Text.builder()
                    .role(Role.USER.getRole())
                    .content(question)
                    .build();
            history.add(text);
            log.trace("==> send message == \n\t{}", text);
            // 用户可以自定义模型参数
            ws.send(XfUtil.genReqMessage(
                    apiAuth, apiVersion, customParams == null ? DEFAULT_PARAMS : customParams,
                    history, functions));
            synchronized (this) {
                // 等待 Listener 线程唤醒(执行完毕)
                this.wait();
            }
            log.trace("==> close connection == {}\n", resp.getMessage());
            ws.close(1000, "close");
            return resp;
        }
    }

}
